# 这是一个示例 Python 脚本。
import torch
from torch import nn, optim

from sort.data.DataHandler import DataHandler, Trainer
from sort.nerve.NeuralNetwork import NeuralNetwork

# 按 ⌃R 执行或将其替换为您的代码。
# 按 双击 ⇧ 在所有地方搜索类、文件、工具窗口、操作和设置。


if __name__ == '__main__':
    data_dir = 'I:\\Workspace\\Dataset\\data_inception\\train'
    batch_size = 8
    # device = "cpu"
    device = (
        "cuda"
        if torch.cuda.is_available()
        else "mps"
        if torch.backends.mps.is_available()
        else "cpu"
    )
    print(device)

    # 加载数据集
    data_handler = DataHandler(data_dir, batch_size)
    train_dataset, val_dataset = data_handler.preprocess_data()
    train_loader, val_loader = data_handler.preprocess_data_loader()

    # 根据数据集确定类 和 类数量
    classes = train_dataset.dataset.classes
    num_classes = len(classes)

    # 创建神经网络
    model = NeuralNetwork(num_classes).to(device)

    # 创建损失函数
    loss_function = nn.CrossEntropyLoss()
    # 创建优化器
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 训练结果查看 & 反向传播和优化
    trainer = Trainer(model, device, loss_function, optimizer)

    # 训练批次
    epochs = 5
    for t in range(epochs):
        print(f"Epoch {t + 1}\n-------------------------------")
        trainer.train(train_loader)
        trainer.test(val_loader)

    model_path = '../../model/ImageSort.pth'

    # 保存模型结构和权重
    # torch.save(model, model_path)

    # 或者只保存模型权重
    torch.save(model.state_dict(), model_path)

    print("Finished Training!")
